#include "kpextBaseCompute.h"

using namespace kpext;

BaseCompute::BaseCompute(kpextVkContext * vkCtx) :
    m_vkCtx(vkCtx),
    m_cols(0),
    m_rows(0),
    m_inBuf(nullptr),
    m_outBuf(nullptr),
    m_postBarrier(true),
    m_cmdBarrier(),
    m_descrPool(nullptr),
    m_descrSetlayout(nullptr),
    m_descrSet(nullptr),
    m_pipeline(nullptr),
    m_workCount({1, 1, 1})
{
    if(!m_vkCtx){
        throw std::runtime_error("Construct kpext::BaseCompute with nullptr !");
        //        return;
    }
}

BaseCompute::~BaseCompute()
{
    if(m_descrPool && m_descrSetlayout){
        m_descrPool->destoryDescriptorSetLayout(m_descrSetlayout);
    }
    m_descrSetlayout = nullptr;
    m_descrSet = nullptr;
    m_descrPool = nullptr;
    m_pipeline = nullptr;
}

void BaseCompute::rebuild(const std::vector<uint32_t> &spirv,
                          const std::vector<std::shared_ptr<BaseBuffer>> & bufs,
                          const std::vector<uint32_t> & specConsts,
                          const std::vector<uint32_t> & pushConsts,
                          const std::array<uint32_t, 3> & workCount
                          )
{
    if(m_descrPool && m_descrSetlayout){
        m_descrPool->destoryDescriptorSetLayout(m_descrSetlayout);
        m_descrSetlayout = nullptr;
    }

    m_descrPool = nullptr;
    {
        m_descrPool = m_vkCtx->descriptorPool(
                    {static_cast<uint32_t>(bufs.size())}
                    );
    }

    if(!m_descrPool){
        SPDLOG_WARN("Create kpext::DescriptorPool failed !");
    }

    m_descrSet = nullptr;

    std::vector<IGetDecriptorInfo *> descrinfos;
    if(bufs.size() < 1){
        SPDLOG_WARN(__FUNCTION__ " input with no storage buffer !");
    }else if(m_descrPool){
        m_cmdBarrier.prepare(
                    bufs,
                    vk::PipelineStageFlagBits::eComputeShader,
                    vk::PipelineStageFlagBits::eComputeShader,
                    vk::AccessFlagBits::eShaderWrite,
                    vk::AccessFlagBits::eShaderRead
                    );
        for (auto & b : bufs) {
            descrinfos.push_back(b->descrInfoGetter());
        }
        m_descrSetlayout = m_descrPool->createDescriptorSetLayout(
                    descrinfos
                    );
    }

    if(m_descrPool && m_descrSetlayout){
        m_descrSet = m_descrPool->recreateDescSet(*m_descrSetlayout);
    }

    if(m_descrSet && descrinfos.size() > 0){
        m_descrPool->writeDescriptorSet(m_descrSet, descrinfos);
    }

    m_pipeline = m_vkCtx->computePipeline();
    if(!m_pipeline){
        SPDLOG_WARN("Create kpext::ComputePipeline failed !");
    }else{
        m_pipeline->m_specConsts = nullptr;
        m_pipeline->m_pushConsts = nullptr;
        if(specConsts.size() > 0){
            m_pipeline->m_specConsts
                    = std::make_shared<kpext::PipelineConstSZ>(
                        specConsts
                        );
        }
        if(pushConsts.size() > 0){
            m_pipeline->m_pushConsts
                    = std::make_shared<kpext::PipelineConstSZ>(
                        pushConsts
                        );
        }
        if(m_descrSetlayout){
            m_pipeline->m_descrSetLayouts.push_back(m_descrSetlayout);
        }
        if(!m_pipeline->rebuild(
                    spirv.size() * sizeof (uint32_t), spirv.data())
                ){
            SPDLOG_WARN("Rebuild kpext::ComputePipeline failed !");
        }
    }

    m_workCount = workCount;
}

void BaseCompute::rebuild(const std::vector<std::shared_ptr<BaseBuffer> > &bufs)
{
    if(!m_descrPool){
        SPDLOG_WARN(__FUNCTION__ " DescriptorPool is invalid!");
        return;
    }

    if(!m_descrSetlayout){
        SPDLOG_WARN(__FUNCTION__ " DescriptorSetLayout is invalid!");
        return;
    }

    uint32_t count = static_cast<uint32_t>(bufs.size());
    if(count < 1){
        SPDLOG_WARN(__FUNCTION__ " kpext::BaseBuffer vector size is zero!");
    }

    for (uint32_t i = 0; i < count; ++i) {
        if(!bufs[i] || !bufs[i]->descrInfoGetter()){
            SPDLOG_ERROR(__FUNCTION__ " input with null pointer!");
            return;
        }
    }

    for (uint32_t i = 0; i < count; ++i) {
        bufs[i]->descrInfoGetter()->setDescBinding(i);
        bufs[i]->descrInfoGetter()->setDescStage(vk::ShaderStageFlagBits::eCompute);
    }

    m_descrPool->resetDescriptorPool();
    //  重新分配描述符集
    m_descrSet = m_descrPool->recreateDescSet(*m_descrSetlayout);
    //  写入绑定信息到描述符集
    std::vector<IGetDecriptorInfo *> descrinfos;
    for(size_t i = 0; i < bufs.size(); ++i){
        descrinfos.push_back(bufs[i]->descrInfoGetter());
    }
    m_descrPool->writeDescriptorSet(m_descrSet,
                                    descrinfos
                                    );
}

void BaseCompute::enablePostBarrier(bool en)
{
    m_postBarrier = en;
}

void BaseCompute::record(const vk::CommandBuffer &commandBuffer)
{
    if(m_pipeline){
        if(m_pipeline->m_pipeline){
            commandBuffer.bindPipeline(
                        vk::PipelineBindPoint::eCompute,
                        *m_pipeline->m_pipeline);
        }


        if(m_pipeline->m_layout){
            commandBuffer.bindDescriptorSets(
                        vk::PipelineBindPoint::eCompute,
                        *m_pipeline->m_layout,
                        0,
                        1,
                        m_descrSet.get(),
                        0,
                        nullptr
                        );

            if(m_pipeline->m_pushConsts){
                commandBuffer.pushConstants(
                            *m_pipeline->m_layout,
                            vk::ShaderStageFlagBits::eCompute,
                            0,
                            static_cast<uint32_t>(m_pipeline->m_pushConsts->size()),
                            m_pipeline->m_pushConsts->data()
                            );
            }
        }

        commandBuffer.dispatch(
                    m_workCount[0],
                m_workCount[1],
                m_workCount[2]
                );

        if(m_postBarrier && m_descrPool){
            m_cmdBarrier.record(commandBuffer);
        }
    }
}
